"""
PoC script to demonstrate PostgreSQL (AGE) / SQL recursive traversal.
Usage: set PG_DSN env or change DSN below.
"""
import os
import json
import time
import psycopg2
from psycopg2.extensions import adapt

from cause_inference.pg import connect_to_postgres, query_topo_entities, query_subgraph, query_cross_host_edges_detail

DSN = os.environ.get('PG_DSN') or 'postgresql://postgres:postgres@localhost:5432/spider'

def setup(conn):
    cur = conn.cursor()
    try:
        cur.execute('CREATE EXTENSION IF NOT EXISTS "age"')
    except Exception:
        # ignore if cannot enable extension (may require superuser)
        pass
    conn.commit()
    # create tables
    cur.execute('''
    CREATE TABLE IF NOT EXISTS observe_entities (
        entity_id text,
        ts bigint,
        type text,
        level text,
        machine_id text,
        attrs jsonb,
        PRIMARY KEY(entity_id, ts)
    )
    ''')
    cur.execute('''
    CREATE TABLE IF NOT EXISTS edges (
        id bigserial PRIMARY KEY,
        from_id text,
        to_id text,
        type text,
        layer text,
        ts bigint,
        attrs jsonb
    )
    ''')
    cur.execute('''CREATE TABLE IF NOT EXISTS timestamps (ts bigint PRIMARY KEY)''')
    conn.commit()

def insert_sample(conn, ts):
    cur = conn.cursor()
    cur.execute("INSERT INTO timestamps (ts) VALUES (%s) ON CONFLICT DO NOTHING", (ts,))
    # insert nodes
    nodes = [
        ('host1', ts, 'host', 'INFO', 'host1', json.dumps({'machine_id': 'host1'})),
        ('proc1', ts, 'process', 'INFO', 'host1', json.dumps({'machine_id': 'host1'})),
        ('disk1', ts, 'disk', 'INFO', 'host1', json.dumps({'machine_id': 'host1'})),
        ('host2', ts, 'host', 'INFO', 'host2', json.dumps({'machine_id': 'host2'})),
        ('proc2', ts, 'process', 'INFO', 'host2', json.dumps({'machine_id': 'host2'})),
    ]
    for n in nodes:
        cur.execute("INSERT INTO observe_entities (entity_id, ts, type, level, machine_id, attrs) VALUES (%s,%s,%s,%s,%s,%s) ON CONFLICT (entity_id, ts) DO UPDATE SET attrs = EXCLUDED.attrs", n)
    # insert edges
    edges = [
        ('proc1', 'host1', 'runs_on', 'layer', ts, json.dumps({})),
        ('disk1', 'proc1', 'belongs_to', 'layer', ts, json.dumps({})),
        ('proc2', 'host2', 'runs_on', 'layer', ts, json.dumps({})),
        ('proc1', 'proc2', 'store_in', 'layer', ts, json.dumps({})),
    ]
    for e in edges:
        cur.execute("INSERT INTO edges (from_id, to_id, type, layer, ts, attrs) VALUES (%s,%s,%s,%s,%s,%s)", e)
    conn.commit()
    # optionally add into AGE graph if available
    USE_AGE = os.environ.get('USE_AGE', 'false').lower() in ['1', 'true', 'yes']
    if not USE_AGE:
        raise Exception('PoC requires Apache AGE enabled. Set USE_AGE=true and ensure extension is installed')
    if USE_AGE:
        try:
            # ensure graph exists
            cur.execute("SELECT create_graph('topo')")
        except Exception:
            pass
        # insert vertices
        for n in nodes:
            entity_id = n[0]
            node_type = n[2]
            machine_id = n[4]
            attrs = n[5]
            entity_q = adapt(entity_id).getquoted().decode()
            type_q = adapt(node_type).getquoted().decode()
            machine_q = adapt(machine_id).getquoted().decode()
            attrs_q = adapt(attrs).getquoted().decode()
            cy = f"CREATE (:ObserveEntity {{entity_id: {entity_q}, ts: {int(ts)}, type: {type_q}, machine_id: {machine_q}, attrs: {attrs_q}}} )"
            try:
                cur.execute("SELECT * FROM cypher(%s, %s)", ('topo', cy))
            except Exception:
                pass
        # insert edges into AGE
        for e in edges:
            from_id, to_id, etype, _, _, _ = e
            from_q = adapt(from_id).getquoted().decode()
            to_q = adapt(to_id).getquoted().decode()
            etype_q = adapt(etype).getquoted().decode()
            cy = f"MATCH (a {{entity_id: {from_q}, ts: {int(ts)}}}), (b {{entity_id: {to_q}, ts: {int(ts)}}}) CREATE (a)-[e:Relation {{ts: {int(ts)}, type: {etype_q}}}]->(b)"
            try:
                cur.execute("SELECT * FROM cypher(%s, %s)", ('topo', cy))
            except Exception:
                pass
        conn.commit()
    # verify AGE installed
    if USE_AGE:
        cur.execute("SELECT count(*) FROM pg_extension WHERE extname = 'age'")
        cnt = cur.fetchone()[0]
        if cnt == 0:
            raise Exception('Apache AGE extension not installed on this server; PoC requires AGE')

def main():
    conn = connect_to_postgres(DSN)
    setup(conn)
    ts = int(time.time())
    insert_sample(conn, ts)
    print('=== nodes ===')
    nodes = query_topo_entities(conn, ts, query_options={'type': 'process'})
    for n in nodes:
        print('node', n)
    print('=== subgraph from proc1 ===')
    nodes_map, edges_map = query_subgraph(conn, ts, 'proc1', ['runs_on', 'belongs_to', 'store_in'], depth=2)
    print('nodes:', list(nodes_map.keys()))
    print('edges:', list(edges_map.keys()))
    print('=== cross host edges ===')
    res = query_cross_host_edges_detail(conn, 'store_in', ts)
    for e in res:
        print('edge', e)

if __name__ == '__main__':
    main()
